import torch
import torch.nn as nn
import onnx

# 1️⃣ 定义一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(4, 2)  # 输入4维 -> 输出2维

    def forward(self, x):
        return self.fc(x)

# 2️⃣ 创建模型并设置为评估模式
model = SimpleNet()
model.eval()

# 3️⃣ 生成一个虚拟输入（假装有一个4维向量输入）
dummy_input = torch.randn(1, 4)

# 4️⃣ 导出模型为 ONNX 格式
torch.onnx.export(
    model,
    dummy_input,
    "simple_net.onnx",          # 导出文件名
    input_names=["input"],      # 输入节点名
    output_names=["output"],    # 输出节点名
    opset_version=11            # ONNX 操作集版本
)

print("✅ 模型已成功导出为 simple_net.onnx")
